// Copyright 2020 Antrea Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package networkpolicy

import (
	"slices"
	"sort"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"antrea.io/antrea/pkg/agent/types"
)

var (
	p110  = types.Priority{TierPriority: 1, PolicyPriority: 1, RulePriority: 0}
	p1120 = types.Priority{TierPriority: 1, PolicyPriority: 1.2, RulePriority: 0}
	p1121 = types.Priority{TierPriority: 1, PolicyPriority: 1.2, RulePriority: 1}
	p1130 = types.Priority{TierPriority: 1, PolicyPriority: 1.3, RulePriority: 0}
	p1131 = types.Priority{TierPriority: 1, PolicyPriority: 1.3, RulePriority: 1}
	p1132 = types.Priority{TierPriority: 1, PolicyPriority: 1.3, RulePriority: 2}
	p1133 = types.Priority{TierPriority: 1, PolicyPriority: 1.3, RulePriority: 3}
	p1140 = types.Priority{TierPriority: 1, PolicyPriority: 1.4, RulePriority: 0}
	p1141 = types.Priority{TierPriority: 1, PolicyPriority: 1.4, RulePriority: 1}
	p1142 = types.Priority{TierPriority: 1, PolicyPriority: 1.4, RulePriority: 2}
	p190  = types.Priority{TierPriority: 1, PolicyPriority: 9, RulePriority: 0}
	p191  = types.Priority{TierPriority: 1, PolicyPriority: 9, RulePriority: 1}
	p192  = types.Priority{TierPriority: 1, PolicyPriority: 9, RulePriority: 2}
	p193  = types.Priority{TierPriority: 1, PolicyPriority: 9, RulePriority: 3}
)

func TestUpdatePriorityAssignment(t *testing.T) {
	tests := []struct {
		name                string
		argsPriorities      []types.Priority
		argsOFPriorities    []uint16
		expectedPriorityMap map[types.Priority]uint16
		expectedOFMap       map[uint16]types.Priority
		expectedSorted      types.ByPriority
	}{
		{
			"in-order",
			[]types.Priority{p110, p1120, p1121},
			[]uint16{10000, 9999, 9998},
			map[types.Priority]uint16{p110: 10000, p1120: 9999, p1121: 9998},
			map[uint16]types.Priority{10000: p110, 9999: p1120, 9998: p1121},
			[]types.Priority{p1121, p1120, p110},
		},
		{
			"reverse-order",
			[]types.Priority{p1121, p1120, p110},
			[]uint16{9998, 9999, 10000},
			map[types.Priority]uint16{p110: 10000, p1120: 9999, p1121: 9998},
			map[uint16]types.Priority{10000: p110, 9999: p1120, 9998: p1121},
			[]types.Priority{p1121, p1120, p110},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			pa := newPriorityAssigner(false)
			for i := 0; i < len(tt.argsPriorities); i++ {
				pa.updatePriorityAssignment(tt.argsOFPriorities[i], tt.argsPriorities[i])
			}
			assert.Equalf(t, tt.expectedPriorityMap, pa.priorityMap, "Got unexpected priorityMap")
			assert.Equalf(t, tt.expectedOFMap, pa.ofPriorityMap, "Got unexpected ofPriorityMap")
			assert.Equalf(t, tt.expectedSorted, pa.sortedPriorities, "Got unexpected sortedPriorities")
		})
	}
}

func TestReassignBoundaryPriorities(t *testing.T) {
	prioritiesToRegister := []types.Priority{p1133, p1132, p1131, p1130}
	tests := []struct {
		name                 string
		lowerBound           uint16
		upperBound           uint16
		originalPriorities   []types.Priority
		originalOfPriorities []uint16
		expectedPriorityMap  map[types.Priority]uint16
		expectedUpdates      map[types.Priority]*priorityUpdate
	}{
		{
			"push-down-single",
			10000,
			10001,
			[]types.Priority{p1140, p1121, p1120},
			[]uint16{10000, 10001, 10002},
			map[types.Priority]uint16{
				p1140: 9996, p1133: 9997, p1132: 9998, p1131: 9999, p1130: 10000, p1121: 10001, p1120: 10002},
			map[types.Priority]*priorityUpdate{p1140: {10000, 9996}},
		},
		{
			"push-down-multiple",
			10000,
			10001,
			[]types.Priority{p190, p1140, p1121, p1120},
			[]uint16{9998, 10000, 10001, 10002},
			map[types.Priority]uint16{
				p190: 9995, p1140: 9996, p1133: 9997, p1132: 9998,
				p1131: 9999, p1130: 10000, p1121: 10001, p1120: 10002},
			map[types.Priority]*priorityUpdate{
				p1140: {10000, 9996},
				p190:  {9998, 9995},
			},
		},
		{
			"push-up-single",
			10000,
			10002,
			[]types.Priority{p1142, p1141, p1140, p1121},
			[]uint16{9998, 9999, 10000, 10002},
			map[types.Priority]uint16{
				p1142: 9998, p1141: 9999, p1140: 10000, p1133: 10001,
				p1132: 10002, p1131: 10003, p1130: 10004, p1121: 10005},
			map[types.Priority]*priorityUpdate{p1121: {10002, 10005}},
		},
		{
			"push-up-multiple",
			10000,
			10002,
			[]types.Priority{p1142, p1141, p1140, p1121, p1120},
			[]uint16{9998, 9999, 10000, 10002, 10003},
			map[types.Priority]uint16{
				p1142: 9998, p1141: 9999, p1140: 10000, p1133: 10001,
				p1132: 10002, p1131: 10003, p1130: 10004, p1121: 10005, p1120: 10006},
			map[types.Priority]*priorityUpdate{
				p1121: {10002, 10005},
				p1120: {10003, 10006},
			},
		},
		{
			"reassign-minimum-possible",
			10000,
			10002,
			[]types.Priority{p193, p192, p191, p190, p1140, p1121, p1120},
			[]uint16{9994, 9995, 9996, 9997, 10000, 10002, 10003},
			map[types.Priority]uint16{
				p193: 9994, p192: 9995, p191: 9996, p190: 9997, p1140: 10000,
				p1133: 10001, p1132: 10002, p1131: 10003, p1130: 10004,
				p1121: 10005, p1120: 10006},
			map[types.Priority]*priorityUpdate{
				p1121: {10002, 10005},
				p1120: {10003, 10006},
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			pa := newPriorityAssigner(false)
			for i, p := range tt.originalOfPriorities {
				pa.updatePriorityAssignment(p, tt.originalPriorities[i])
			}
			priorityUpdates := map[types.Priority]*priorityUpdate{}
			err := pa.reassignBoundaryPriorities(tt.lowerBound, tt.upperBound, prioritiesToRegister, priorityUpdates)
			assert.NoError(t, err, "Error occurred in reassignment")
			assert.Equalf(t, tt.expectedPriorityMap, pa.priorityMap, "priorityMap unexpected after reassignment")
			assert.Equalf(t, tt.expectedUpdates, priorityUpdates, "priority updates unexpected after reassignment")
		})
	}
}

func TestInsertConsecutivePriorities(t *testing.T) {
	prioritiesToRegister := []types.Priority{p1133, p1132, p1131, p1130}
	pa := newPriorityAssigner(false)
	insertionLow := pa.initialOFPriority(prioritiesToRegister[0])
	insertionHigh := pa.initialOFPriority(prioritiesToRegister[3])
	tests := []struct {
		name                  string
		originalPriorities    []types.Priority
		originalOfPriorities  []uint16
		expectedOFPriorityMap map[uint16]types.Priority
	}{
		{
			"empty-map",
			[]types.Priority{},
			[]uint16{},
			map[uint16]types.Priority{
				insertionLow: p1133, insertionLow + 1: p1132, insertionLow + 2: p1131, insertionHigh: p1130,
			},
		},
		{
			"irrelevant-high-priority",
			[]types.Priority{p110},
			[]uint16{insertionLow + 100},
			map[uint16]types.Priority{
				insertionLow: p1133, insertionLow + 1: p1132, insertionLow + 2: p1131, insertionHigh: p1130,
				insertionLow + 100: p110,
			},
		},
		{
			"irrelevant-low-priority",
			[]types.Priority{p1140},
			[]uint16{insertionLow - 100},
			map[uint16]types.Priority{
				insertionLow - 100: p1140,
				insertionLow:       p1133, insertionLow + 1: p1132, insertionLow + 2: p1131, insertionHigh: p1130,
			},
		},
		{
			"irrelevant-surrounding-priorities",
			[]types.Priority{p1141, p1140, p1121, p1120},
			[]uint16{insertionLow - 100, insertionLow - 99, insertionLow + 99, insertionLow + 100},
			map[uint16]types.Priority{
				insertionLow - 100: p1141, insertionLow - 99: p1140,
				insertionLow: p1133, insertionLow + 1: p1132, insertionLow + 2: p1131, insertionHigh: p1130,
				insertionLow + 99: p1121, insertionLow + 100: p1120,
			},
		},
		{
			"overlapping-low-priorities",
			[]types.Priority{p1141, p1140},
			[]uint16{insertionLow + 1, insertionLow + 2},
			map[uint16]types.Priority{
				insertionLow + 1: p1141, insertionLow + 2: p1140,
				insertionLow + 3 + zoneOffset: p1133, insertionLow + 4 + zoneOffset: p1132,
				insertionLow + 5 + zoneOffset: p1131, insertionLow + 6 + zoneOffset: p1130,
			},
		},
		{
			"overlapping-high-priorities",
			[]types.Priority{p1121, p1120},
			[]uint16{insertionLow + 1, insertionLow + 2},
			map[uint16]types.Priority{
				insertionLow - zoneOffset - 3: p1133, insertionLow - zoneOffset - 2: p1132,
				insertionLow - zoneOffset - 1: p1131, insertionLow - zoneOffset: p1130,
				insertionLow + 1: p1121, insertionLow + 2: p1120,
			},
		},
		{
			"priorities-with-small-gap",
			[]types.Priority{p1141, p1140, p1121, p1120},
			[]uint16{insertionLow + 1, insertionLow + 2, insertionLow + 9, insertionLow + 10},
			map[uint16]types.Priority{
				insertionLow + 1: p1141, insertionLow + 2: p1140,
				insertionLow + 4: p1133, insertionLow + 5: p1132, insertionLow + 6: p1131, insertionLow + 7: p1130,
				insertionLow + 9: p1121, insertionLow + 10: p1120,
			},
		},
		{
			"gap-just-enough-for-insertion",
			[]types.Priority{p1141, p1140, p1121, p1120},
			[]uint16{insertionLow - 1, insertionLow, insertionLow + 5, insertionLow + 6},
			map[uint16]types.Priority{
				insertionLow - 1: p1141, insertionLow: p1140,
				insertionLow + 1: p1133, insertionLow + 2: p1132, insertionLow + 3: p1131, insertionLow + 4: p1130,
				insertionLow + 5: p1121, insertionLow + 6: p1120,
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			pa := newPriorityAssigner(false)
			for i, p := range tt.originalOfPriorities {
				pa.updatePriorityAssignment(p, tt.originalPriorities[i])
			}
			priorityUpdates := map[types.Priority]*priorityUpdate{}
			err := pa.insertConsecutivePriorities(prioritiesToRegister, priorityUpdates)
			assert.NoError(t, err, "Error occurred in priority insertion")
			assert.Equalf(t, tt.expectedOFPriorityMap, pa.ofPriorityMap, "priorityMap unexpected after insertion")
		})
	}
}

func TestRegisterPrioritiesAndRevert(t *testing.T) {
	pa := newPriorityAssigner(false)
	prioritiesToRegister := []types.Priority{p1132, p1131, p1130, p190, p191}
	insertionPoint1132 := pa.initialOFPriority(p1132)
	insertionPoint191 := pa.initialOFPriority(p191)
	pa.updatePriorityAssignment(insertionPoint1132-1, p1140)
	pa.updatePriorityAssignment(insertionPoint1132+2, p1121)
	pa.updatePriorityAssignment(insertionPoint1132+3, p1120)

	expectedOFMapAfterRegister := map[uint16]types.Priority{
		insertionPoint1132 - 2: p1140,
		insertionPoint1132 - 1: p1132, insertionPoint1132: p1131, insertionPoint1132 + 1: p1130,
		insertionPoint1132 + 2: p1121, insertionPoint1132 + 3: p1120,
		insertionPoint191: p191, insertionPoint191 + 1: p190,
	}
	_, revertFunc, err := pa.registerPriorities(prioritiesToRegister)
	assert.NoError(t, err, "Error occurred in priority registration")
	assert.Equalf(t, expectedOFMapAfterRegister, pa.ofPriorityMap, "priorityMap unexpected after registration")

	expectedOFMapAfterRevert := map[uint16]types.Priority{
		insertionPoint1132 - 1: p1140, insertionPoint1132 + 2: p1121, insertionPoint1132 + 3: p1120,
	}
	revertFunc()
	assert.Equalf(t, expectedOFMapAfterRevert, pa.ofPriorityMap, "priorityMap unexpected after revert")
}

func TestRegisterDuplicatePriorities(t *testing.T) {
	pa1 := newPriorityAssigner(false)
	pa2 := newPriorityAssigner(false)
	prioritiesToRegister := []types.Priority{p1131, p1130}
	prioritiesToRegisterDuplicate := []types.Priority{p1130, p1131, p1130, p1130, p1130, p1131, p1130}

	_, _, err := pa1.registerPriorities(prioritiesToRegister)
	assert.NoError(t, err, "Error occurred in priority registration")
	_, _, err2 := pa2.registerPriorities(prioritiesToRegisterDuplicate)
	assert.NoError(t, err2, "Error occurred in priority registration")
	ofPriority1130, _ := pa1.getOFPriority(p1130)
	ofPriority1130Dup, _ := pa2.getOFPriority(p1130)
	assert.Equal(t, ofPriority1130, ofPriority1130Dup)
	ofPriority1131, _ := pa1.getOFPriority(p1131)
	ofPriority1131Dup, _ := pa2.getOFPriority(p1131)
	assert.Equal(t, ofPriority1131, ofPriority1131Dup)
}

func TestRegister500PrioritiesNoOverflow(t *testing.T) {
	tests := []struct {
		name         string
		generateFunc func() []types.Priority
	}{
		{
			"500-default-tier-priorities",
			func() []types.Priority {
				var priorities []types.Priority
				for i := 1; i <= 500; i++ {
					priority := types.Priority{
						TierPriority:   defaultTierPriority, // 250
						PolicyPriority: float64(i),
						RulePriority:   0,
					}
					priorities = append(priorities, priority)
				}
				return priorities
			},
		},
		{
			"500-tier1-small-policy-priorities",
			func() []types.Priority {
				var priorities []types.Priority
				for i := 1; i <= 500; i++ {
					priority := types.Priority{
						TierPriority:   1,
						PolicyPriority: float64(i) * 0.01,
						RulePriority:   0,
					}
					priorities = append(priorities, priority)
				}
				return priorities
			},
		},
		{
			"500-tier20-with-rule-priorities",
			func() []types.Priority {
				var priorities []types.Priority
				for i := 1; i <= 500; i++ {
					priority := types.Priority{
						TierPriority:   20,
						PolicyPriority: float64(i) * 5,
						RulePriority:   int32(i),
					}
					priorities = append(priorities, priority)
				}
				return priorities
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			pa := newPriorityAssigner(false)

			// Generate priorities for this test case
			priorities := tt.generateFunc()
			// Make a copy since registerPriorities modifies the slice order
			originalPriorities := make([]types.Priority, len(priorities))
			copy(originalPriorities, priorities)
			sort.Sort(types.ByPriority(originalPriorities))

			_, _, err := pa.registerPriorities(priorities)
			require.NoError(t, err)

			var sortedOFPriorities []uint16
			for _, priority := range originalPriorities {
				ofPriority, exists := pa.getOFPriority(priority)
				require.True(t, exists, "Priority %v should be registered", priority)
				sortedOFPriorities = append(sortedOFPriorities, ofPriority)
			}

			// Verify that OpenFlow priorities are sorted in ascending order (higher precedence gets lower OpenFlow priority)
			require.True(t, slices.IsSorted(sortedOFPriorities), "OpenFlow priorities should be sorted in ascending order (higher precedence gets lower OpenFlow priority)")
		})
	}
}

func generatePriorities(tierPriority, start, end int32, policyPriority float64) []types.Priority {
	priorities := make([]types.Priority, end-start+1)
	for i := start; i <= end; i++ {
		priorities[i-start] = types.Priority{TierPriority: tierPriority, PolicyPriority: policyPriority, RulePriority: i - start}
	}
	return priorities
}

func TestRegisterAllOFPriorities(t *testing.T) {
	pa := newPriorityAssigner(true)
	maxPriorities := generatePriorities(253, int32(baselinePolicyBottomPriority), int32(baselinePolicyTopPriority), 5)
	_, _, err := pa.registerPriorities(maxPriorities)
	assert.NoError(t, err, "Error occurred in registering max number of allowed priorities in baseline tier")

	// Check that all priorities registered in baseline tier are sorted in ofPriorities space
	sortedMaxPriorities := make([]types.Priority, len(maxPriorities))
	copy(sortedMaxPriorities, maxPriorities)
	sort.Sort(types.ByPriority(sortedMaxPriorities))

	var baselineOFPriorities []uint16
	for _, priority := range sortedMaxPriorities {
		ofPriority, exists := pa.getOFPriority(priority)
		require.True(t, exists, "Priority %v should be registered", priority)
		baselineOFPriorities = append(baselineOFPriorities, ofPriority)
	}
	// Verify that OpenFlow priorities are sorted in ascending order (higher precedence gets lower OpenFlow priority)
	require.True(t, slices.IsSorted(baselineOFPriorities), "Baseline tier OpenFlow priorities should be sorted in ascending order")

	extraPriority := types.Priority{
		TierPriority:   253,
		PolicyPriority: 5,
		RulePriority:   int32(baselinePolicyTopPriority) - int32(baselinePolicyBottomPriority) + 1,
	}
	_, _, err = pa.registerPriorities([]types.Priority{extraPriority})
	assert.Errorf(t, err, "Error should be raised after max number of priorities are registered in baseline tier")

	pa = newPriorityAssigner(false)
	consecPriorities1 := generatePriorities(5, int32(policyBottomPriority), 10000, 5)
	_, _, err = pa.registerPriorities(consecPriorities1)

	assert.NoError(t, err, "Error occurred before registering max number of allowed priorities")
	consecPriorities2 := generatePriorities(10, 10001, int32(policyTopPriority), 5)
	_, _, err = pa.registerPriorities(consecPriorities2)
	assert.NoError(t, err, "Error occurred in registering max number of allowed priorities")

	// Create a combined sorted copy of all registered priorities
	allRegisteredPriorities := make([]types.Priority, 0, len(consecPriorities1)+len(consecPriorities2))
	allRegisteredPriorities = append(allRegisteredPriorities, consecPriorities1...)
	allRegisteredPriorities = append(allRegisteredPriorities, consecPriorities2...)
	sort.Sort(types.ByPriority(allRegisteredPriorities))

	// Get OpenFlow priorities for all registered priorities in sorted order
	var allOFPriorities []uint16
	for _, priority := range allRegisteredPriorities {
		ofPriority, exists := pa.getOFPriority(priority)
		require.True(t, exists, "Priority %v should be registered", priority)
		allOFPriorities = append(allOFPriorities, ofPriority)
	}
	// Verify that OpenFlow priorities are sorted in ascending order (higher precedence gets lower OpenFlow priority)
	require.True(t, slices.IsSorted(allOFPriorities), "All registered OpenFlow priorities should be sorted in ascending order")

	_, _, err = pa.registerPriorities([]types.Priority{extraPriority})
	assert.Errorf(t, err, "Error should be raised after max number of priorities are registered")
}
